import os

import torch
import pygame

from math import *
from modules.Particles import Particles


def calculate_viewing_distance(fov, screen_width):
    d = (screen_width / 2.0) / tan(fov / 2.0)
    return d


def move_points(points,dv):
    for p in points:
        p.update(p + dv)
    points.sort(key=lambda p:p.z, reverse=True)

def rotate_points(points,p_float,v):
    for p in points:
        p.update(p.rotate(p_float,v))
    points.sort(key=lambda p:p.z, reverse=True)

def calculate_quantile(x, p):
    sorted_x = torch.sort(x)[0]
    length = len(sorted_x)
    index = torch.floor(p * length / 100)
    return sorted_x[index]

def main():
    import pygame
    from pygame.math import Vector3

    pygame.init()

    # 让pygame完全控制鼠标
    # os.environ['SDL_VIDEO_WINDOW_POS'] = "%d,%d" % (10, 10)
    pygame.mouse.set_visible(False)
    pygame.event.set_grab(True)
    fov = 90.  # 视野
    rov = 5
    mov = 90
    FPS = 60
    is_grab = True
    # 设置屏幕大小和标题
    screen_width = 800
    screen_height = 600
    SCREEN_SIZE = (screen_width,screen_height)
    screen = pygame.display.set_mode((screen_width, screen_height))
    pygame.display.set_caption("Particle Motion")
    font = pygame.font.SysFont('simhei', 24)
    clock = pygame.time.Clock()

    viewing_distance = calculate_viewing_distance(radians(fov), screen_width)
    camera_position = Vector3(0.0, 0.0, 0.0)  # 摄像机位置
    center_x, center_y = SCREEN_SIZE
    center_x /= 2
    center_y /= 2

    ball_w, ball_h = 1, 1
    ball_center_x = ball_w / 2
    ball_center_y = ball_h / 2

    # 设置颜色
    white = (255, 255, 255)
    black = (0, 0, 0)

    # 生成粒子位置和速度的数组

    num_particles = int(1e4)
    device = "cuda"
    init_time = 0
    # mass = 1
    time_step = 3600*24
    # particles = []
    current_time = init_time

    mass_scale = 1e24
    min_mass,max_mass = int(1e5),int(1e6)
    radius_scale = 1e9
    positions = torch.randint(0, screen_width, (num_particles,3), dtype=torch.float32,device=torch.device(device)) * radius_scale
    masses = torch.randint(min_mass,max_mass,(num_particles,1), dtype=torch.float32,device=torch.device(device))
    particles = Particles(num_particles,positions, masses, init_time, mass_scale,device=device)

    # 主循环
    running = True
    while running:
        points = []
        print("current_time：{}".format(current_time))
        # 处理事件
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                running = False
        pressed_keys = pygame.key.get_pressed()

        if is_grab:
            mouse_rel = pygame.mouse.get_rel()
            drx = mouse_rel[0] * 2 * rov
            dry = mouse_rel[1] * 2 * rov
        else:
            drx, dry = 0, 0
        dx, dy, dz = 0.0, 0.0, 0.0

        # 旋转
        ## 绕y轴旋转
        rotate_points(points,drx * 1 / FPS, Vector3(0.0, -1.0, 0.0))
        ## 绕x轴旋转
        rotate_points(points,dry * 1 / FPS, Vector3(-1.0, 0.0, 0.0))


        # current_time
        # for i in range(num_particles):
        #     particles[i]
        #     positions[:, i] += velocities[:, i] * dt

            # 清屏
        screen.fill(white)


        # 绘制粒子位置
        # positions = particles.positions / torch.max(particles.positions,dim=0)[0]
        positions = particles.positions / radius_scale
        # if current_time == 0:
        for i in range(num_particles):
            # print(i)
            position = positions[i]
            # print(particles.accelerations[i])
            if torch.isnan(position.sum()):
                pass
            elif torch.isinf(position.sum()):
                pass
            else:
                # position_scaled = (position/radius_scale)
                pygame.draw.circle(screen, black, (int(position[0]),
                                                   int(position[1])), radius=1)

                # points.append(Vector3(position[0]/radius_scale, position[1]/radius_scale, position[2]/radius_scale))
                # pygame.draw.circle(screen, black, (position[0], position[1]), radius=2)

        # move_points(points, Vector3(dx, dy, dz))
        #
        # for point in points:
        #     x, y, z = point - camera_position
        #     if z > 0:
        #         x = x * viewing_distance / z
        #         y = -y * viewing_distance / z
        #         x += center_x
        #         y += center_y
        #         bw = min(100 / point.magnitude() * ball_w, center_x)
        #         # screen.blit(ball, (x - ball_center_x, y - ball_center_y))
        #         print(x)
        #         pygame.draw.circle(
        #             screen,
        #             black,
        #             (int(x - ball_center_x), int(y - ball_center_y)),
        #             int(bw)
        #         )


        ## 更新粒子的位置
        current_time += time_step
        particles.add_time(current_time)
        particles.update_positions()

        # 更新屏幕显示
        # pygame.display.update()
        # time_passed = clock.tick(FPS)
        pygame.display.flip()
        # break


    # 退出pygame
    pygame.quit()


def matplotlib_main():
    import numpy as np
    import matplotlib.pyplot as plt
    import matplotlib.animation as animation
    from matplotlib.animation import FuncAnimation

    # 生成粒子位置和速度的数组

    num_particles = int(3)
    device = "cuda"
    init_time = 0
    # mass = 1
    time_step = 3600 * 10
    trace_len = 50
    # particles = []
    global current_time
    current_time = init_time
    # 设置屏幕大小和标题
    screen_width = 800
    screen_height = 600
    max_radius = 1
    mass_scale = 1e24
    min_mass, max_mass = int(5), int(6)
    radius_scale = 9e9
    show_trace = True

    positions = torch.rand((num_particles, 3), dtype=torch.float32,
                              device=torch.device(device)) * max_radius * radius_scale
    masses = torch.randint(min_mass, max_mass, (num_particles, 1), dtype=torch.float32, device=torch.device(device))
    velocities = torch.randint(0,100,(num_particles, 3),device=torch.device(device))
    particles = Particles(num_particles, positions, masses, velocities,init_time, mass_scale, device=device)

    positions = particles.positions / radius_scale
    # if current_time == 0:
    # print(positions[0])
    global xlines
    global ylines
    global zlines

    xlines = []
    ylines = []
    zlines = []

    # 定义初始图形
    fig = plt.figure()
    ax = fig.add_subplot(111, projection='3d')
    ax.scatter(list(positions[:, 0].cpu()), list(positions[:, 1].cpu()), list(positions[:, 2].cpu()), c='b')
    xlines = positions[:, 0].cpu()
    ylines = positions[:, 1].cpu()
    zlines = positions[:, 2].cpu()

    # for i in range(num_particles):
    #     # print(i)
    #     position = positions[i]
    #     # print(particles.accelerations[i])
    #     if torch.isnan(position.sum()):
    #         pass
    #     elif torch.isinf(position.sum()):
    #         pass
    #     else:
    #         ax.scatter(position[0].cpu(), position[1].cpu(), position[2].cpu())
    ax.set_xlabel('X Label/ly')
    ax.set_ylabel('Y Label/ly')
    ax.set_zlabel('Z Label/ly')
    ax.set_title('N-body 3D Scatter Plot')

    def update(i):
        global xlines
        global ylines
        global zlines
        plt.cla()
        # 更新数据

        global current_time
        current_time += time_step
        particles.add_time(current_time)
        particles.update_positions()

        positions = particles.positions / radius_scale
        ax.scatter(list(positions[:,0].cpu()), list(positions[:,1].cpu()), list(positions[:,2].cpu()), c='b')
        # 绘制轨迹
        xlines = torch.vstack([xlines,positions[:, 0].cpu()])
        ylines = torch.vstack([ylines,positions[:, 1].cpu()])
        zlines = torch.vstack([zlines,positions[:, 2].cpu()])

        if xlines.shape[0] > trace_len:
            xlines = xlines[-trace_len:, ]
            ylines = ylines[-trace_len:, ]
            zlines = zlines[-trace_len:, ]



        # xlines = torch.tensor(positions[:, 0],device=torch.device(device))
        # ylines = torch.tensor(positions[:, 1],device=torch.device(device))
        # zlines = torch.tensor(positions[:, 2],device=torch.device(device))
        ax.set_xlabel('X Label/ly')
        ax.set_ylabel('Y Label/ly')
        ax.set_zlabel('Z Label/ly')
        ax.set_title('N-body 3D Scatter Plot')
        # cmap_list = [plt.cm.get_cmap('RdYlBu')(i) for i in np.linspace(0, 1, 50)]
        if show_trace:
            for i in range(xlines.shape[1]):
                x = xlines[:,i]
                y = ylines[:,i]
                z = zlines[:,i]

                # p_yx = np.polyfit(y, x, 3)
                # x_out = np.polyval(p_yx, y)
                # p_yz = np.polyfit(y, z, 3)
                # z_out = np.polyval(p_yz, y)
                ax.plot3D(x, y, z, color="blue")

        # for i in range(num_particles):
        #     # print(i)
        #     position = positions[i]
            # print(particles.accelerations[i])
            # if torch.isnan(position.sum()):
            #     pass
            # elif torch.isinf(position.sum()):
            #     pass
            # else:
                # ax.scatter(position[0], position[1], position[2])

                # 更新图形坐标
                # ax.cla()
                # ax.scatter(position[0], position[1].cpu(), position[2].cpu(), c='b')
                # ax.scatter(position[0].cpu(), position[1].cpu(), position[2].cpu(),c='b')

        return ax,

    ani = FuncAnimation(fig, update, frames=range(500),interval=50)
    ani.save("./3-body.gif",bitrate=6000,fps=60,dpi=300)
    plt.show()


if __name__ == '__main__':
    # main()

    matplotlib_main()